# -*- coding: utf-8 -*-
"""
PDF 加载器（依赖 BasePdfLoader 通用流程）
仅支持 事故案例 文件的处理方式
"""
import json
import os
from abc import ABC
from pathlib import Path
from typing import Dict

import requests
from loguru import logger
from zai import ZhipuAiClient

try:
    from document_loader.pdf_utils import PDFMetadata, BasePdfLoader, PdfCommonUtils
except ImportError:
    import sys

    CURRENT_DIR = os.path.dirname(__file__)
    PROJECT_ROOT = os.path.abspath(os.path.join(CURRENT_DIR, os.pardir))
    if PROJECT_ROOT not in sys.path:
        sys.path.append(PROJECT_ROOT)
    from document_loader.pdf_utils import PDFMetadata, BasePdfLoader, PdfCommonUtils


class AcPdfLoader(BasePdfLoader):
    """
    事故案例
    """
    def __init__(self, file_path, delete_old_file: bool = False):
        super().__init__(file_path, delete_old_file)

        self.API_URL = "https://api.textin.com/ai/service/v1/pdf_to_markdown"
        self.headers = {
            "x-ti-app-id": "8e5ac65e5c09273415e4d6275493e3c3",
            "x-ti-secret-code": "f240bfbfcebbb42f9be0dafce5c0f44d",
            "Content-Type": "application/octet-stream",
        }
        self.API_KEY = "fabc0f1f271448f9a0b7207ae47f77c5.o8ITXRoUG1N1dTEe"  # 智谱API_KEY
        self.client = ZhipuAiClient(api_key=self.API_KEY)

        self.SUPPORTED_EXT = {".png", ".jpg", ".jpeg", ".pdf", ".bmp", ".tiff", ".webp", ".doc", ".docx", ".html", ".mhtml"}

    def _file_to_markdown(self, json_path: Path, md_path: Path) -> None:
        """调用 TextIn API 将 PDF 转为 Markdown，并存储原始结果 json。"""
        try:
            params = {
                "page_start": 0,
                "page_count": 1000,
                "parse_mode": "scan",
                "table_flavor": "html",
                "apply_document_tree": 1,
                "page_details": 1,
                "markdown_details": 1,
                "get_image": "objects",
                "image_output_type": "default",
                "raw_ocr": 1,
            }
            file_data = self.file_path.read_bytes()
            resp = requests.post(self.API_URL, headers=self.headers, params=params, data=file_data)
            resp.raise_for_status()
            j = resp.json()
            if j.get("code") != 200:
                raise RuntimeError(f"API 错误: {j.get('message')}")
            md_result = j["result"]
            json_path.write_text(json.dumps(md_result, ensure_ascii=False, indent=2), encoding="utf-8")
            md_path.write_text(md_result.get("markdown", ""), encoding="utf-8")
            logger.debug(f"✅ 已保存：{json_path} 和 {md_path}")

        except Exception as e:
            logger.error(f"PDF转换失败: {e}")
            raise

    def _extract_accident_info_with_glm(self, markdown_content, parsed_md_path):
        prompt = f"""
            请从以下事故调查报告中抽取核心信息，按照指定结构返回JSON格式：

            事故调查报告内容：
            {markdown_content}

            请按以下结构抽取信息：

            {{
                "一、事故基本信息": {{
                    "事故标题": "",
                    "事故类型": "",
                    "事故等级": "",
                    "报告日期": "",
                    "涉事单位": "",
                    "事故摘要": "",
                }},
                "二、事故经过与后果": {{
                    "事故经过": "",
                    "人员伤亡情况": "",
                    "财产损失情况": "",
                    "环境影响": "",
                    "作业中断情况": "",

                }},
                "三、直接原因": {{
                    "不安全的行为": "",
                    "不安全的状态": "",
                }},
                "四、根本原因": {{
                    "管理因素": "",
                    "技术因素": "",
                    "人员因素": "",
                    "环境因素": ""
                }},
                "五、事故处理与整改": {{
                    "应急响应措施": "",
                    "事故调查结论": "",
                    "责任追究": "",
                    "整改措施": "",
                    "整改落实情况": "",
                }},
                "六、其他辅助信息": {{
                    "报告来源": "",
                    "报告日期": "",
                    "相关方": "",
                    "关键词/标签": "",
                }}
            }}

            说明：
            1、事故类型：根据事故调查报告内容，判断事故类型，如：火灾、爆炸、中毒、触电、机械伤害、高处坠落、物体打击、车辆伤害、坍塌、其他等。
            2、事故等级：根据事故调查报告内容，判断事故等级，如：一般事故、较大事故、重大事故、特别重大事故。
            3、事故摘要：总结事故经过，如：事故发生时间、事故发生地点、事故发生原因、事故发生后果等。
            4、事故时间：返回格式为 xxxx/xx/xx。

            要求：
            1. 严格按照上述JSON结构返回
            2. 如果某项信息在报告中没有，填入"未提及"
            3. 信息抽取的要准确完整，不要编造内容
            4. 你的任务是抽取并总结事故案例的核心信息，不要编造内容，不要添加任何解释性内容，不要添加任何主观性内容，不要添加任何评价性内容。
        """

        try:
            response = self.client.chat.completions.create(
                model="glm-4.5",
                messages=[
                    {"role": "system", "content": "你是一个专业的事故调查报告分析专家，擅长从文档中抽取结构化信息。"},
                    {"role": "user", "content": prompt}
                ],
                thinking={
                    "type": "enabled",
                },
                response_format={"type": "json_object"}
            )

            result_text = response.choices[0].message.content.strip()
            result = json.loads(result_text)
            return result

        except Exception as e:
            print(f"GLM调用失败 {parsed_md_path}: {e}")
            return None

    def load(self, delete_old_file: bool = None) -> Dict:
        file_pure_name = self.file_path.stem
        md_save_dir = self.base_dir
        md_save_dir.mkdir(parents=True, exist_ok=True)

        parsed_json_path = md_save_dir / f"{file_pure_name}.json"
        parsed_md_path = md_save_dir / f"{file_pure_name}.md"
        extracted_json_path = md_save_dir / f"{file_pure_name}_extracted.json"

        if self.delete_old_file or not parsed_md_path.exists():
            logger.debug(f"[Loader] Generating markdown: {parsed_md_path}")
            self._pdf_to_markdown(parsed_json_path, parsed_md_path)

        # 不存在则解析，存在直接读取结果文件
        if self.delete_old_file or not extracted_json_path.exists():
            logger.debug(f"[Loader] Extracting markdown: {parsed_md_path}")
            with open(parsed_md_path, 'r', encoding='utf-8') as f:
                markdown_content = f.read()
            result = self._extract_accident_info_with_glm(markdown_content, parsed_md_path)
            with open(extracted_json_path, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
        else:
            with open(extracted_json_path, 'r', encoding='utf-8') as f:
                result = json.load(f)

        return result
